import os
import random
import shutil


def train_all(origin_dir, val_ratio=0.1):
    """
    default format of directory：
    ----images
    ----labels
    :param train_ratio: ratio of train set
    :param origin_dir:to store images and labels
    :return:
    """
    # get all prefix of labels
    images_dir = os.path.join(origin_dir, r"Images")
    images_list = os.listdir(images_dir)
    images_prefix = []
    for img in images_list:
        images_prefix.append(os.path.splitext(img)[0])
    random.shuffle(images_prefix)
    # Set up train set number and validation set number
    val_number = int(len(images_prefix) * val_ratio)
    val_list = images_prefix[:val_number]
    # move pictures and annotations to corresponding directory
    if not os.path.exists(os.path.join(origin_dir, r"labels")):
        os.mkdir(os.path.join(origin_dir, r"labels"))
    if not os.path.exists(os.path.join(origin_dir, r"labels", r"train")):
        os.mkdir(os.path.join(origin_dir, r"labels", r"train"))
    if not os.path.exists(os.path.join(origin_dir, r"labels", r"val")):
        os.mkdir(os.path.join(origin_dir, r"labels", r"val"))
    if not os.path.exists(os.path.join(origin_dir, r"images")):
        os.mkdir(os.path.join(origin_dir, r"images"))
    if not os.path.exists(os.path.join(origin_dir, r"images", r"train")):
        os.mkdir(os.path.join(origin_dir, r"images", r"train"))
    if not os.path.exists(os.path.join(origin_dir, r"images", r"val")):
        os.mkdir(os.path.join(origin_dir, r"images", r"val"))
    # 生成训练集
    with open(os.path.join(origin_dir, "train.txt"), "w") as train_txt:
        for train_ele in images_prefix:
            picture_name = train_ele + r".jpg"
            picture_path = r"./images/train/" + picture_name + "\n"
            train_txt.write(picture_path)
            old_picture = os.path.join(origin_dir, r"Images", picture_name)
            new_picture = os.path.join(origin_dir, r"images", r"train", picture_name)
            shutil.copy(old_picture, new_picture)
            anno_name = train_ele + r".txt"
            old_anno = os.path.join(origin_dir, r"Labels", anno_name)
            if os.path.exists(old_anno):
                new_anno = os.path.join(origin_dir, r"labels", r"train", anno_name)
                shutil.copy(old_anno, new_anno)
    # 生成验证集
    with open(os.path.join(origin_dir, "val.txt"), "w") as val_txt:
        for val_ele in val_list:
            picture_name = val_ele + r".jpg"
            picture_path = r"./images/val/" + picture_name + "\n"
            val_txt.write(picture_path)
            old_picture = os.path.join(origin_dir, r"Images", picture_name)
            new_picture = os.path.join(origin_dir, r"images", r"val", picture_name)
            shutil.copy(old_picture, new_picture)
            anno_name = val_ele + r".txt"
            old_anno = os.path.join(origin_dir, r"Labels", anno_name)
            if os.path.exists(old_anno):
                new_anno = os.path.join(origin_dir, r"labels", r"val", anno_name)
                shutil.copy(old_anno, new_anno)


def train_val_split(origin_dir, val_ratio=0.1):
    """
    default format of directory：
    ----images
    ----labels
    :param train_ratio: ratio of train set
    :param origin_dir:to store images and labels
    :return:
    """
    # get all prefix of images
    images_dir = os.path.join(origin_dir, r"Images")
    images_list = os.listdir(images_dir)
    images_prefix = []
    for img in images_list:
        images_prefix.append(os.path.splitext(img)[0])
    random.shuffle(images_prefix)
    # Set up train set number and validation set number
    train_number = int(len(images_prefix) * (1 - val_ratio))
    train_list = images_prefix[:train_number + 1]
    val_list = images_prefix[train_number + 1:]
    # move pictures and annotations to corresponding directory
    if not os.path.exists(os.path.join(origin_dir, r"labels")):
        os.mkdir(os.path.join(origin_dir, r"labels"))
    if not os.path.exists(os.path.join(origin_dir, r"labels", r"train")):
        os.mkdir(os.path.join(origin_dir, r"labels", r"train"))
    if not os.path.exists(os.path.join(origin_dir, r"labels", r"val")):
        os.mkdir(os.path.join(origin_dir, r"labels", r"val"))
        if not os.path.exists(os.path.join(origin_dir, r"images")):
            os.mkdir(os.path.join(origin_dir, r"images"))
    if not os.path.exists(os.path.join(origin_dir, r"images", r"train")):
        os.mkdir(os.path.join(origin_dir, r"images", r"train"))
    if not os.path.exists(os.path.join(origin_dir, r"images", r"val")):
        os.mkdir(os.path.join(origin_dir, r"images", r"val"))
    # 生成训练集
    with open(os.path.join(origin_dir, "train.txt"), "w") as train_txt:
        for train_ele in train_list:
            picture_name = train_ele + r".jpg"
            picture_path = r"./images/train/" + picture_name + "\n"
            train_txt.write(picture_path)
            old_picture = os.path.join(origin_dir, r"Images", picture_name)
            new_picture = os.path.join(origin_dir, r"images", r"train", picture_name)
            shutil.move(old_picture, new_picture)
            anno_name = train_ele + r".txt"
            old_anno = os.path.join(origin_dir, r"Labels", anno_name)
            if os.path.exists(old_anno):
                new_anno = os.path.join(origin_dir, r"labels", r"train", anno_name)
                shutil.move(old_anno, new_anno)
    # 生成验证集
    with open(os.path.join(origin_dir, "val.txt"), "w") as val_txt:
        for val_ele in val_list:
            picture_name = val_ele + r".jpg"
            picture_path = r"./images/val/" + picture_name + "\n"
            val_txt.write(picture_path)
            old_picture = os.path.join(origin_dir, r"Images", picture_name)
            new_picture = os.path.join(origin_dir, r"images", r"val", picture_name)
            shutil.move(old_picture, new_picture)
            anno_name = val_ele + r".txt"
            old_anno = os.path.join(origin_dir, r"Labels", anno_name)
            if os.path.exists(old_anno):
                new_anno = os.path.join(origin_dir, r"labels", r"val", anno_name)
                shutil.move(old_anno, new_anno)


if __name__ == "__main__":
    root_path = r"/home/kaijia/algo-env/datasets/hq_safeguard_dataset/yolov7"
    train_val_split(root_path)
    # train_all(root_path)
